import numpy as np
from .base import BaseTransformer
from .error_checker import PreErrorChecker
from scipy.stats import mode


class StandardScaler(BaseTransformer):
    def __init__(self):
        super().__init__()
        self.mean_ = None  # 均值
        self.std_ = None   # 标准差

    def fit(self, X, y=None):
        """
        计算均值和标准差，保存下来供 transform 使用。
        """
        PreErrorChecker.check_array(X)
        PreErrorChecker.check_missing_values(X)
        PreErrorChecker.check_constant_columns(X)

        self.mean_ = np.mean(X, axis=0)
        self.std_ = np.std(X, axis=0)

    def transform(self, X):
        """
        使用保存的均值和标准差对数据进行标准化处理。
        """
        PreErrorChecker.check_array(X)
        if self.mean_ is None or self.std_ is None:
            raise ValueError("The transformer is not fitted yet. Please call 'fit' first.")

        return (X - self.mean_) / self.std_


class SimpleImputer(BaseTransformer):
    """
    SimpleImputer 类用于填充缺失数据。根据指定的策略（均值、中位数或最频繁值），填充数据中的缺失值。
    """
    def __init__(self, strategy='mean'):
        super().__init__()
        self.strategy = strategy
        self.fill_value_ = None
        self.check_valid_strategy()

    def check_valid_strategy(self):
        """检查填充策略是否有效"""
        if self.strategy not in ['mean', 'median', 'most_frequent']:
            raise ValueError(
                f"Invalid strategy '{self.strategy}'. Valid strategies are ['mean', 'median', 'most_frequent'].")

    def fit(self, X, y=None):
        """
        根据指定策略计算填充值（均值、最频繁值等）。
        """
        PreErrorChecker.check_array(X)
        PreErrorChecker.check_constant_columns(X)

        if self.strategy == 'mean':
            self.fill_value_ = np.nanmean(X, axis=0)
        elif self.strategy == 'median':
            self.fill_value_ = np.nanmedian(X, axis=0)
        elif self.strategy == 'most_frequent':
            self.fill_value_ = mode(X, nan_policy='omit').mode[0]


    def transform(self, X):
        """
        使用填充值填充数据中的缺失值。
        """
        PreErrorChecker.check_array(X)
        return np.where(np.isnan(X), self.fill_value_, X)


class MinMaxScaler(BaseTransformer):
    """
    MinMaxScaler 类用于将数据缩放到指定的最小和最大值之间，通常是 [0, 1]。
    """
    def __init__(self, feature_range=(0, 1)):
        super().__init__()
        self.feature_range = feature_range
        self.min_ = None
        self.max_ = None

    def fit(self, X, y=None):
        """
        计算最小值和最大值，保存下来供 transform 使用。
        """
        PreErrorChecker.check_array(X)
        PreErrorChecker.check_missing_values(X)
        PreErrorChecker.check_constant_columns(X)

        self.min_ = np.min(X, axis=0)
        self.max_ = np.max(X, axis=0)

    def transform(self, X):
        """
        将数据缩放到指定的范围。
        """
        PreErrorChecker.check_array(X)
        if self.min_ is None or self.max_ is None:
            raise ValueError("The transformer is not fitted yet. Please call 'fit' first.")

        return (X - self.min_) / (self.max_ - self.min_) * (self.feature_range[1] - self.feature_range[0]) + self.feature_range[0]
